#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <sys/un.h>
#include <errno.h>
#include <unistd.h>
#include <fcntl.h>

#include <stdarg.h>
#include <string.h>

#include <glib.h>

#include "util.h"
#include "command-client.h"
#include "nautilus-nutstore.h"

/**
 * Connection Attempt
 */
typedef struct {
    NutstoreCommandClient   *client;
    guint                   num_attempts;
} ConnectionAttempt;

typedef struct {
    NutstoreCommandClientConnectionAttemptHook hook;
    gpointer user_data;
} NutstoreCommandClientConnectionAttempt;

/**
 * create command 
 */
static NutstoreCommand *_command_new_va(NutstoreCommandRespHandler handler,
        gboolean b_add_idle,
        gpointer user_data, 
        const char *command,
        va_list ap)
{
    NutstoreCommand *cmd;
    gchar *key;

    cmd = g_new(NutstoreCommand, 1);

    cmd->command_name = g_strdup(command);
    cmd->command_args = g_hash_table_new_full((GHashFunc) g_str_hash,
            (GEqualFunc) g_str_equal,
            (GDestroyNotify) g_free,
            (GDestroyNotify) g_strfreev);

    cmd->cancelled = FALSE;
    cmd->handler = handler;
    cmd->b_add_idle = b_add_idle;
    cmd->user_data = user_data;
    cmd->response = NULL;

    while ((key = va_arg(ap, char *)) != NULL) {
        gchar **values;

        values = g_new(gchar *, 2);
        values[0] = g_strdup(va_arg(ap, char *));
        values[1] = NULL;

        g_hash_table_insert(cmd->command_args, g_strdup(key), values);
    }

    return cmd;
}

NutstoreCommand * nutstore_command_new(NutstoreCommandRespHandler handler, 
        gboolean b_add_idle,
        gpointer user_data,
        const char *command, ...)
{
    NutstoreCommand *cmd;
    va_list ap;

    va_start(ap, command);
    cmd = _command_new_va(handler, b_add_idle, user_data, command, ap);
    va_end(ap);

    return cmd;
}

static void nutstore_command_free(NutstoreCommand *cmd)
{
    if (cmd->response != NULL) {
        g_hash_table_unref(cmd->response);
    }

    g_free(cmd->command_name);
    if (cmd->command_args != NULL) {
        g_hash_table_unref(cmd->command_args);
    }
    g_free(cmd);
}

NutstoreCommand * nutstore_command_new_simple(const char *cmd_name)
{
    NutstoreCommand *cmd;

    cmd = g_new(NutstoreCommand, 1);

    cmd->command_name = g_strdup(cmd_name);
    cmd->command_args = NULL;
    cmd->cancelled = FALSE;
    cmd->handler = NULL;
    cmd->b_add_idle = FALSE;
    cmd->user_data = NULL;
    cmd->response = NULL;

    return cmd;
}

/**
 * Call all onconnect hooks
 */
static gboolean on_connect(NutstoreCommandClient *client) 
{
    g_hook_list_invoke(&(client->onconnect_hooklist), FALSE);
    return FALSE;
}

/**
 * Call all ondisconnect hooks
 */
static gboolean on_disconnect(NutstoreCommandClient *client) 
{
    g_hook_list_invoke(&(client->ondisconnect_hooklist), FALSE);
    return FALSE;
}

/**
 * Call all connection attempt hooks
 */
static gboolean on_connection_attempt(ConnectionAttempt *ca) 
{
    GList *ll;

    for (ll = ca->client->conn_attempt_hooklist; ll != NULL; ll = g_list_next(ll)) {
        NutstoreCommandClientConnectionAttempt *clientca =
            (NutstoreCommandClientConnectionAttempt *)(ll->data);
        clientca->hook(ca->num_attempts, clientca->user_data);
    }

    g_free(ca);

    return FALSE;
}

/**
 * receiv all arguments
 */
#define MAX_NUM_ARGS    (16)
static gboolean receive_args_until_done(GIOChannel *chan, 
        GHashTable *return_table,
        GError **err) 
{
    GIOStatus iostat;
    GError *tmp_error = NULL;
    guint numargs = 0;

    while (1) {
        gchar *line;
        gsize term_pos;

        /* if we are getting too many args, connection could be malicious */
        if (numargs >= MAX_NUM_ARGS) {
            g_set_error(err, g_quark_from_static_string("malicious connection"),
                    0, "malicious connection");
            return FALSE;
        }

        /* get the string */
        iostat = g_io_channel_read_line(chan, &line, NULL, &term_pos, &tmp_error);
        if (iostat == G_IO_STATUS_ERROR || tmp_error != NULL) {
            g_free(line);
            if (tmp_error != NULL) {
                g_propagate_error(err, tmp_error);
            }
            return FALSE;
        } else if (iostat == G_IO_STATUS_EOF) {
            g_free(line);
            g_set_error(err, g_quark_from_static_string("connection closed"),
                    0, "connection closed");
            return FALSE;
        }

        // All response MUST end with "done\n"
        if (strcmp(line, "done\n") == 0) {
            g_free(line);
            break;
        } else {
            gboolean parse_result;

            // replace the line end '\n' with '\0' 
            line[term_pos] = '\0';

            parse_result = nutstore_util_command_parse_arg(line, return_table);
            g_free(line);

            if (FALSE == parse_result) {
                //g_printf("command_parse_arg() failed\n");
                g_set_error(err, g_quark_from_static_string("response parse error"),
                        0, "response parse error");
                return FALSE;
            }
        }

        numargs += 1;
    }

    return TRUE;
}

/**
 * Custom functions to get all keys of hash table
 * because g_hash_table_get_keys() is added in GLib 2.14.
 */
#if !GLIB_CHECK_VERSION(2, 14, 0)
static void my_g_hash_table_get_keys_helper(gpointer key,
        gpointer value,
        GList **ud) 
{
    *ud = g_list_append(*ud, key);
}
#endif

static GList *my_g_hash_table_get_keys(GHashTable *ght) {
    GList *list = NULL;

#if !GLIB_CHECK_VERSION(2, 14, 0)
    g_hash_table_foreach(ght, (GHFunc)my_g_hash_table_get_keys_helper, &list);
#else
    list = g_hash_table_get_keys(ght);
#endif

    return list;
}

/*
   sends a command to the nutstore server
   returns an hash of the return values

   in theory, this should disconnection errors
   but it doesn't matter right now, any error is a sufficient
   condition to disconnect
 */
static GHashTable *send_command_to_daemon(GIOChannel *chan, 
        const gchar *command_name,
        GHashTable *args, 
        GError **err) 
{
    GError *tmp_error = NULL;
    GIOStatus iostat;
    gsize bytes_trans;
    gchar *line;

    g_assert(chan != NULL);
    g_assert(command_name != NULL);


#define WRITE_OR_DIE_SANI(s,l) {                        \
    gchar *sani_s;                                      \
    sani_s = nutstore_util_sanitize(s);                 \
    iostat = g_io_channel_write_chars(chan, sani_s, l,  \
            &bytes_trans, &tmp_error);                  \
    g_free(sani_s);                                     \
    if (iostat == G_IO_STATUS_ERROR || iostat == G_IO_STATUS_AGAIN) {   \
        if (tmp_error != NULL) {                        \
            g_propagate_error(err, tmp_error);          \
        }                                               \
        return NULL;                                    \
    }                                                   \
}

#define WRITE_OR_DIE(s,l) {                             \
    iostat = g_io_channel_write_chars(chan, s, l,       \
            &bytes_trans, &tmp_error);                  \
    if (iostat == G_IO_STATUS_ERROR || tmp_error != NULL) {             \
        if (tmp_error != NULL) {                        \
            g_propagate_error(err, tmp_error);          \
        }                                               \
        return NULL;                                    \
    }                                                   \
}

    /* send command to server */
    WRITE_OR_DIE_SANI(command_name, -1);
    WRITE_OR_DIE("\n", -1);

    if (args != NULL) {
        GList *keys, *li;

        /* g_hash_table_get_keys in Glib < 2.14 has some bug? */
        keys =  my_g_hash_table_get_keys(args);

        for (li = keys; li != NULL; li = g_list_next(li)) {
            int i;
            gchar **value;

            WRITE_OR_DIE_SANI((gchar *) li->data, -1);

            value = g_hash_table_lookup(args, li->data);
            for (i = 0; value[i] != NULL; i++) {
                WRITE_OR_DIE("\t", -1);
                WRITE_OR_DIE_SANI(value[i], -1);
            }
            WRITE_OR_DIE("\n", -1);
        }

        g_list_free(keys);
    }

    WRITE_OR_DIE("done\n", -1);

#undef WRITE_OR_DIE
#undef WRITE_OR_DIE_SANI

    g_io_channel_flush(chan, &tmp_error);
    if (tmp_error != NULL) {
        //g_printf("g_io_channel_flush() failed\n");
        g_propagate_error(err, tmp_error);
        return NULL;
    }

    /* now we have to read the data */
    iostat = g_io_channel_read_line(chan, &line, NULL, NULL, &tmp_error);
    if (iostat == G_IO_STATUS_ERROR) {
        g_assert(line == NULL);
        g_propagate_error(err, tmp_error);
        return NULL;
    } else if (iostat == G_IO_STATUS_AGAIN) {
        g_assert(line == NULL);
        g_set_error(err, 
                g_quark_from_static_string("command connection timed out"),
                0,
                "command connection timed out");
        return NULL;
    } else if (iostat == G_IO_STATUS_EOF) {
        g_assert(line == NULL);
        g_set_error(err,
                g_quark_from_static_string("command connection closed"),
                0,
                "command connection closed");
        return NULL;
    }

    //g_printf("read_line() success; line=%s", line);

    /* All response MUST start with line: rsp\n */
    if (strncmp(line, "rsp\n", 4) == 0) {
        GHashTable *return_table = 
        g_hash_table_new_full((GHashFunc) g_str_hash,
                (GEqualFunc) g_str_equal,
                (GDestroyNotify) g_free,
                (GDestroyNotify) g_strfreev);

        g_free(line);
        line = NULL;

        receive_args_until_done(chan, return_table, &tmp_error);
        if (tmp_error != NULL) {
            //g_printf("receive_args_until_done() failed\n");
            g_hash_table_destroy(return_table);
            g_propagate_error(err, tmp_error);
            return NULL;
        }

        return return_table;
    }  else {   /* otherwise */
        /* read errors off until we get done */
        while (strncmp(line, "done\n", 5) != 0) {
            g_free(line);
            line = NULL;

            /* clear string */
            iostat = g_io_channel_read_line(chan, &line, NULL,
                    NULL, &tmp_error);
            if (iostat == G_IO_STATUS_ERROR) {
                g_assert(line == NULL);
                g_propagate_error(err, tmp_error);
                return NULL;
            } else if (iostat == G_IO_STATUS_AGAIN) {
                g_assert(line == NULL);
                g_set_error(err,
                        g_quark_from_static_string("command connection timed out"),
                        0,
                        "command connection timed out");
                return NULL;
            } else if (iostat == G_IO_STATUS_EOF) {
                g_assert(line == NULL);
                g_set_error(err,
                        g_quark_from_static_string("command connection closed"),
                        0,
                        "command connection closed");
                return NULL;
            }

            /* we got our line */
        } 

        g_free(line);
        return NULL;
    }
}

/** 
 * execute the command handler in the gui thread. After the handler is
 * invoked, the command is freed. So, the command handler should never free
 *  the command 
 */
static gboolean finish_nutstore_command_async(NutstoreCommand *cmd)
{
    g_assert(cmd->handler != NULL);
    g_assert(cmd->b_add_idle);

    cmd->handler(cmd);
    nutstore_command_free(cmd);

    // MUST return FALSE to remove from the main loop.
    return FALSE;
}

static void finish_nutstore_command(NutstoreCommand *cmd) 
{
    if (cmd->handler != NULL) {
        if (cmd->b_add_idle) {
            g_idle_add((GSourceFunc)finish_nutstore_command_async, cmd);
        } else {
            cmd->handler(cmd);
            nutstore_command_free(cmd);
        }
    } else {
        nutstore_command_free(cmd);
    }
}

static void process_command(GIOChannel *chan, 
        NutstoreCommand *cmd,
        GError **gerr) 
{
    GError *tmp_gerr = NULL;
    GHashTable *response;

    /* send status command to server */
    response = send_command_to_daemon(chan, cmd->command_name,
            cmd->command_args, &tmp_gerr);
    if (tmp_gerr != NULL) {
        g_assert(response == NULL);
        g_propagate_error(gerr, tmp_gerr);
        return;
    }

    /* great, the server did the command perfectly,
     * now call the handler with the response 
     */
    cmd->response = response;  
    finish_nutstore_command(cmd);

    return;
}

static gboolean check_connection(GIOChannel *chan) 
{
    gchar fake_buf[4096];
    gsize bytes_read;
    GError *tmp_error = NULL;
    GIOFlags flags;
    GIOStatus ret, iostat;

    flags = g_io_channel_get_flags(chan);

    /* set non-blocking */
    ret = g_io_channel_set_flags(chan, flags | G_IO_FLAG_NONBLOCK, NULL);
    if (ret == G_IO_STATUS_ERROR) {
        return FALSE;
    }

    iostat = g_io_channel_read_chars(chan, fake_buf,
            sizeof(fake_buf), 
            &bytes_read, &tmp_error);

    ret = g_io_channel_set_flags(chan, flags, NULL);
    if (ret == G_IO_STATUS_ERROR) {
        return FALSE;
    }

    /* this makes us disconnect from bad servers
     * (those that send us information without us asking for it) 
     */
    return iostat == G_IO_STATUS_AGAIN;
}

static gpointer nutstore_command_client_thread(NutstoreCommandClient *data);

static void end_request(NutstoreCommand *cmd) 
{
    if ((gpointer (*)(NutstoreCommandClient *data))cmd != &nutstore_command_client_thread) {
        cmd->response = NULL;
        finish_nutstore_command(cmd);
    } else {
        /** This a special command to reset the connection. Do not free it */
    }
}

static gboolean try_connect(int *pSocket)
{
    int sock = -1;
    struct sockaddr_in addr;
    struct timeval tv = { 3, 0 };   // Timeout: 3 seconds
    int saved_flags;
    int ret;

    /* intialize address structure */
    addr.sin_family = AF_INET;                                              
    addr.sin_port = htons(COMMAND_SERVER_PORT);
    inet_aton("127.0.0.1", &addr.sin_addr);

    sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (sock < 0) {
        goto TRY_CONNECT_FAIL;
    }

    /* set timeout on socket, to protect against bad servers */
    ret = setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(struct timeval));
    if (ret < 0) {
        debug("setsockeopt() RCVTIMEO failed");
        goto TRY_CONNECT_FAIL;
    }

    ret = setsockopt(sock, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(struct timeval));
    if (ret < 0) {
        debug("setsockeopt() SNDTIMEO failed");
        goto TRY_CONNECT_FAIL;
    }

    /* save current flag for later restore */ 
    saved_flags = fcntl(sock, F_GETFL, 0);
    if (saved_flags < 0) {
        debug("fcntl() GET failed");
        goto TRY_CONNECT_FAIL;
    }
            
    /* set native non-blocking, for connect timeout */
    ret = fcntl(sock, F_SETFL, saved_flags | O_NONBLOCK);
    if (ret < 0) {
        debug("fcntl() SET failed");
        goto TRY_CONNECT_FAIL;
    }

    /* if there was an error we have to try again later */
    if (connect(sock, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
        if (errno == EINPROGRESS) {
            fd_set writers;
            struct timeval tv = { 1, 0 };

            FD_ZERO(&writers);
            FD_SET(sock, &writers);

            /* if nothing was ready after 3 seconds, fail out homie */
            if (select(sock+1, NULL, &writers, NULL, &tv) == 0) {
                debug("connection timeout");
                goto TRY_CONNECT_FAIL;
            }

            // Re-try
            if (connect(sock, (struct sockaddr *) &addr, sizeof(addr)) < 0) {
                debug("couldn't connect to server after 1 second"); 
                goto TRY_CONNECT_FAIL;
            }
        } else {    /* errno != EINPROGRESS */
            debug("bad connection");
            goto TRY_CONNECT_FAIL;
        }
    }

    /* set back to blocking */
    if (fcntl(sock, F_SETFL, saved_flags) < 0) {
        debug("fcntl() restore failed");
        goto TRY_CONNECT_FAIL;
    }

    *pSocket = sock;
    return TRUE;

TRY_CONNECT_FAIL:
    if (sock >= 0) {
        close(sock);
    }
    return FALSE;
}

static gpointer nutstore_command_client_thread(NutstoreCommandClient *client) 
{
    int num_attempts = 1;

    while (1) {
        int sock;
        GIOChannel *chan = NULL;
        GError *gerr = NULL;

        if (try_connect(&sock) == FALSE) {
            ConnectionAttempt *ca = g_new(ConnectionAttempt, 1);
            ca->client = client;
            ca->num_attempts = num_attempts;

            g_idle_add((GSourceFunc) on_connection_attempt, ca);

            // wait 2 seconds before retrying again
            g_usleep(2 * G_USEC_PER_SEC);
            num_attempts++;
            continue;
        } else {
            num_attempts = 0;
        }

        /* connected */
        debug("command client connected");

        chan = g_io_channel_unix_new(sock);

        // set close_on_unref
        g_io_channel_set_close_on_unref(chan, TRUE);
        g_io_channel_set_line_term(chan, "\n", -1);

#define SET_CONNECTED_STATE(s)     {    \
    g_mutex_lock(client->mutex);	    \
    client->is_connected = s;			\
    g_mutex_unlock(client->mutex);	    \
}

        SET_CONNECTED_STATE(TRUE);

        g_idle_add((GSourceFunc) on_connect, client);

        while (1) {
            NutstoreCommand *cmd;

            while (1) {
                GTimeVal gtv;

                g_get_current_time(&gtv);

                // We do not need to check the connection aggresively, because
                // we can know whether it is broken when recv the response from the peer
                g_time_val_add(&gtv, 5 * G_USEC_PER_SEC);

                /* get a request from nautilus */
                cmd = g_async_queue_timed_pop(client->cmd_queue, &gtv);
                if (cmd != NULL) {
                    break;
                } else {
                    if (check_connection(chan) == FALSE) {
                        goto BADCONNECTION;
                    }
                }
            }

            /* special command to reset the connection: this pointer should be unique */
            if ((gpointer (*)(NutstoreCommandClient *data))cmd == &nutstore_command_client_thread) {
                debug("got a reset request");
                goto BADCONNECTION;
            }

            debug("processing command %s", cmd->command_name);
            process_command(chan, cmd, &gerr);
            debug("done.");

            if (gerr != NULL) {
                //	debug("COMMAND ERROR*****************************");
                /* mark this request as never to be completed */
                end_request(cmd);

                debug("command error: %s", gerr->message);

                g_error_free(gerr);

BADCONNECTION:
                /* grab all the rest of the data off the async queue and mark it
                   never to be completed, who knows how long we'll be disconnected */
                while ((cmd = g_async_queue_try_pop(client->cmd_queue)) != NULL) {
                    end_request(cmd);
                }

                g_io_channel_unref(chan);

                SET_CONNECTED_STATE(FALSE);

                /* call the disconnect handler */
                g_idle_add((GSourceFunc) on_disconnect, client);

                break;
            }
        }

#undef SET_CONNECTED_STATE
}

return NULL;
}

/* thread safe */
gboolean
nutstore_command_client_is_connected(NutstoreCommandClient *client) 
{
    gboolean is_connected;

    g_mutex_lock(client->mutex);
    is_connected = client->is_connected;
    g_mutex_unlock(client->mutex);

    return is_connected;
}

/* thread safe */
void nutstore_command_client_force_reconnect(NutstoreCommandClient *client) 
{
    if (nutstore_command_client_is_connected(client) == TRUE) {
        debug("forcing command to reconnect");
        nutstore_command_client_request(client, (NutstoreCommand *)&nutstore_command_client_thread);
    }
}

/* thread safe */
void
nutstore_command_client_request(NutstoreCommandClient *client, 
        NutstoreCommand *cmd) 
{
    g_async_queue_push(client->cmd_queue, cmd);
}

/* should only be called once on initialization */
void
nutstore_command_client_setup(NutstoreCommandClient *client) 
{
    client->cmd_queue = g_async_queue_new();
    client->mutex = g_mutex_new();
    client->is_connected = FALSE;
    client->conn_attempt_hooklist = NULL;

    g_hook_list_init(&(client->onconnect_hooklist), sizeof(GHook));
    g_hook_list_init(&(client->ondisconnect_hooklist), sizeof(GHook));
}

void
nutstore_command_client_add_on_disconnect_hook(NutstoreCommandClient *client,
        NutstoreCommandClientConnectHook func,
        gpointer user_data) 
{
    GHook *newhook;

    newhook = g_hook_alloc(&(client->ondisconnect_hooklist));
    newhook->func = func;
    newhook->data = user_data;

    g_hook_append(&(client->ondisconnect_hooklist), newhook);
}

void
nutstore_command_client_add_on_connect_hook(NutstoreCommandClient *client,
        NutstoreCommandClientConnectHook func,
        gpointer user_data) 
{
    GHook *newhook;

    newhook = g_hook_alloc(&(client->onconnect_hooklist));
    newhook->func = func;
    newhook->data = user_data;

    g_hook_append(&(client->onconnect_hooklist), newhook);
}

void
nutstore_command_client_add_connection_attempt_hook(NutstoreCommandClient *client,
        NutstoreCommandClientConnectionAttemptHook hook,
        gpointer user_data) 
{
    NutstoreCommandClientConnectionAttempt *newhook;

    debug("shouldn't be here...");

    newhook = g_new(NutstoreCommandClientConnectionAttempt, 1);
    newhook->hook = hook;
    newhook->user_data = user_data;

    client->conn_attempt_hooklist = g_list_append(client->conn_attempt_hooklist, newhook);
}

/* should only be called once on initialization */
void
nutstore_command_client_start(NutstoreCommandClient *client) {
    /* setup the connect to the command server */
    debug("starting command thread");
    g_thread_create((gpointer (*)(gpointer data)) nutstore_command_client_thread,
            client, FALSE, NULL);
}

/* thread safe */
void nutstore_command_client_send_simple_command(NutstoreCommandClient *client, 
        const char *command) 
{
    NutstoreCommand *cmd = nutstore_command_new_simple(command);

    nutstore_command_client_request(client, cmd);
}

/* thread safe */
/* this is the C API, there is another send_command_to_daemon()
   that is more the actual over the wire command */
void nutstore_command_client_send_command(NutstoreCommandClient *client, 
        NutstoreCommandRespHandler handler,
        gboolean b_add_idle,
        gpointer user_data,
        const char *command, ...) 
{
    NutstoreCommand *cmd;
    va_list ap;

    va_start(ap, command);
    cmd = _command_new_va(handler, b_add_idle, user_data, command, ap);
    va_end(ap);

    nutstore_command_client_request(client, cmd);
}
